package top.yudoge.dogeeditor.llm.rewrite

import com.google.gson.Gson
import dev.langchain4j.data.message.ChatMessage
import dev.langchain4j.data.message.SystemMessage
import dev.langchain4j.data.message.UserMessage
import dev.langchain4j.model.chat.response.ChatResponse
import dev.langchain4j.model.chat.response.StreamingChatResponseHandler
import dev.langchain4j.model.embedding.EmbeddingModel
import dev.langchain4j.store.embedding.EmbeddingSearchRequest
import org.springframework.beans.factory.annotation.Autowired
import org.springframework.stereotype.Component
import top.yudoge.dogeeditor.app.utils.ChatHistoryUtils
import top.yudoge.dogeeditor.app.utils.Message
import top.yudoge.dogeeditor.commons.utils.convertList
import top.yudoge.dogeeditor.commons.utils.toJsonStr
import top.yudoge.dogeeditor.db.api.DataSourceMetaDataApi
import top.yudoge.dogeeditor.llm.embedding.EmbeddingStoreHolder
import top.yudoge.dogeeditor.llm.memory.SessionMemory
import top.yudoge.dogeeditor.llm.objects.mro.TableMRO
import java.util.LinkedList

@Component
class SQLGenerate2ChainNode : TaskChainNode {

    @Autowired
    private lateinit var embeddingStoreHolder: EmbeddingStoreHolder

    @Autowired
    private lateinit var dataSourceMetaDataApi: DataSourceMetaDataApi

    @Autowired
    private lateinit var sessionMemory: SessionMemory

    @Autowired
    private lateinit var embeddingModel: EmbeddingModel

    override fun process(ctx: ChainContext) {
        doProcess(ctx)
    }

    fun doProcess(ctx: ChainContext, lastErrorMessage: String? = null) {
        sendReasoningCallback(ctx, "Generating SQL...")

        val messages = buildMessages(ctx, lastErrorMessage)

        val streamingModel = ctx.streamingModel
        val callback = ctx.callback

        streamingModel.chat(messages, object : StreamingChatResponseHandler {
            override fun onPartialResponse(p0: String?) {
                if (p0 != null) {
                    callback.token(p0)
                }
            }

            override fun onCompleteResponse(p0: ChatResponse?) {
                p0?.let {
                    val aiMessage = it.aiMessage().text()
                    ChatHistoryUtils.addChatHistory(ctx.query.sessionId, Message("ai", aiMessage), sessionMemory)

//                    val SQL = extractSQLStatement(aiMessage);
//                    if (SQL != null) {
//                        putChainStore(ctx, GENERATED_SQL, SQL)
//                    }
                }
                callback.complete()
//                if (ctx.query.selfChecking) {
//                    emitNext(ctx)
//                } else {
//                    callback.complete()
//                }
            }

            override fun onError(p0: Throwable?) {
                callback.error(p0)
            }
        })

    }

    private fun extractSQLStatement(aiMessage: String): String? {
        val lines = mutableListOf<String>()
        var startFounded = false
        for (line in aiMessage.lines()) {
            if (line.lowercase() == "```sql" && !startFounded) {
                startFounded = true
                continue
            }
            if (line == "```") {
                break
            }
            if (startFounded) {
                lines.add(line)
            }
        }
        return lines.joinToString("\n")
    }


    private fun buildMessages(ctx: ChainContext, lastErrorMessage: String? = null): List<ChatMessage> {

        val userCode = ctx.query.code
        val userTask = ctx.query.userTask

        val dataSource = dataSourceMetaDataApi.getDataSources().first { it.identifier == ctx.query.dataSourceIdentifier }!!


        val relatedTableNames = mutableSetOf<String>()
        relatedTableNames.addAll(getChainStore(ctx, GetPossibleTableChainNode.POSSIBLE_TABLES) as? List<String> ?: ArrayList())
        relatedTableNames.addAll(ctx.query.tableSelected)
        val relatedTableInfo = ArrayList<TableMRO>()

        for (tableName in relatedTableNames) {
            var columns =
                dataSourceMetaDataApi.getColumns(ctx.query.dataSourceIdentifier, ctx.query.databaseName, ctx.query.schemaName, tableName)
            if (!columns.isNullOrEmpty()) {
                var tableMRO = TableMRO(tableName, columns.convertList())
                relatedTableInfo.add(tableMRO)
            }
        }
        sessionMemory.put(ctx.query.sessionId, HIS_TABLES_KEY, relatedTableNames.toJsonStr())

        var store = embeddingStoreHolder.getStore("use_knowledge")

        var embed = embeddingModel.embed(userTask).content()
        var result = store.search(EmbeddingSearchRequest(embed, 10, 0.5, null))
        var historyKnowledge = result.matches().map { it -> it.embedded().text() }.toJsonStr()
        var prompt = loadPrompt(
            "/rewrite2.template", mapOf(
                "dbms" to dataSource.vendor.vendorName,
                "user_task" to userTask,
                "user_code" to userCode,
                "knowledge" to historyKnowledge,
                "tables" to relatedTableInfo.toJsonStr(),
                "history" to ChatHistoryUtils.listChatHistory(ctx.query.sessionId, sessionMemory).toJsonStr()
            )
        )

        println(prompt)

        ChatHistoryUtils.addChatHistory(ctx.query.sessionId, Message("user", userTask), sessionMemory)

        return buildList {
            add(UserMessage.from(prompt))
        }

    }



    companion object {
        const val GENERATED_SQL = "generated_sql"
        const val HIS_TABLES_KEY = "history_tables"
    }


}